'''
@Company: TWL
@Author: xue jian
@Email: xuejian@kanzhun.com
@Date: 2020-04-09 19:10:56
'''
# -*- coding: UTF-8 -*-
import numpy as np
import tensorflow as tf
from base_train import BaseTrain

class CompareTrain(BaseTrain):

    def __init__(self, file_path, dates, watch_num, batch_size, thread_num, read_data_parallel, all_tf=False):
        super(CompareTrain, self).__init__(file_path, dates, watch_num, batch_size, thread_num,
                                              read_data_parallel)
        self.all_tf = all_tf
    def get_train_object(self):
        print('get_train_object')
        return []


    def init_all_model(self, compare_trains):
        # self.sess = tf.Session(config=tf.ConfigProto(device_count={"CPU": self.thread_num}, log_device_placement=False))
#        sess = tf.Session(config=tf.ConfigProto(device_count={"CPU": self.thread_num}, log_device_placement=False))
        sess = tf.Session()
        for compare_train in compare_trains:
            compare_train.get_feature_placeholder()
            compare_train.make_feature_column()
            compare_train.define_network()
            compare_train.sess = sess
        sess.run(tf.global_variables_initializer())
        sess.run(tf.tables_initializer())

    def consume_data(self, data_queue, aucs, loss_f):
        compare_trains = self.get_train_object()
        if self.all_tf:
            self.init_all_model(compare_trains)
        else:
            for compare_train in compare_trains:
                compare_train.init_model()
        batch_data = []
        label = []
        pred = []
        losses = []
        for i in range(len(compare_trains)):
            pred.append([])
            losses.append([])
        while True:
            tmp = data_queue.get()
            if len(tmp) > 0:
                batch_data.append(tmp)
            if len(batch_data) >= self.batch_size:
                batch_data = np.asarray(batch_data)
                batch_l = batch_data[:, self.fea_index['deal_type']].astype(np.int16)
                for j, compare_train in enumerate(compare_trains):
                    tmp_pred, tmp_loss = compare_train.batch_train(batch_data[:self.batch_size], batch_l)
                    pred[j].extend(tmp_pred)
                    losses[j].append(tmp_loss)
                batch_data = []
                label.extend(batch_l)
                if len(label) >= self.watch_num:
                    tmp_auc = self.cal_aucs(label, pred)
                    aucs.put(tmp_auc)
                    print('auc = ', tmp_auc)
                    put_loss = []
                    for k in range(len(compare_trains)):
                        put_loss.append(np.asarray(losses[k]).mean())
                    loss_f.put(put_loss)
                    label = []
                    for k in range(len(compare_trains)):
                        pred[k] = []
                        losses[k] = []

    def single_train(self):
        def read_date_data(d_f, n):
            all_d = []
            k = 0
            while k < n:
                line = d_f.readline()
                if line == "":
                    break
                line = line[:-1].split("\t")
                tmp_label = self.make_label(line[self.fea_index['deal_type']])
                if tmp_label == -99:
                    continue
                tmp_line = []
                for j, l in enumerate(line):
                    if j == self.fea_index['deal_type']:
                        tmp_line.append(tmp_label)
                    elif l in ['NULL', '\\N']:
                        tmp_line.append(-2)
                    else:
                        tmp_line.append(l)
                all_d.append(tmp_line)
                k += 1
            return all_d

        compare_trains = self.get_train_object()
        if self.all_tf:
            self.init_all_model(compare_trains)
        else:
            for compare_train in compare_trains:
                compare_train.init_model()
        left_data = []
        for date in self.dates:
            date_paths = self.get_train_file(date)
            for date_path in date_paths:
                with open(date_path, 'r') as date_file:
                    all_data = read_date_data(date_file, self.watch_num - len(left_data))
                    all_data = left_data + all_data
                    while len(all_data) >= self.watch_num:
                        all_pred = []
                        all_label = []
                        all_loss = []
                        for i in range(len(compare_trains)):
                            all_pred.append([])
                            all_loss.append([])
                        all_data = np.asarray(all_data)
                        for i in range(len(all_data) // self.batch_size):
                            train_data = all_data[i * self.batch_size: (i + 1) * self.batch_size]
                            batch_l = train_data[:, self.fea_index['deal_type']].astype(np.int16)
                            for j, compare_train in enumerate(compare_trains):
                                tmp_pred, tmp_loss = compare_train.batch_train(train_data, batch_l)
                                all_pred[j].extend(tmp_pred)
                                all_loss[j].append(tmp_loss)
    
                            all_label.extend(batch_l)
    
                        auc = self.cal_aucs(np.asarray(all_label), all_pred)
                        print("auc = ", auc)
                        self.aucs.append(auc)
                        watch_loss = []
                        for los in all_loss:
                            watch_loss.append(np.asarray(los).mean())
                        print("loss = ", watch_loss)
                        self.losses.append(watch_loss)
                        all_data = read_date_data(date_file, self.watch_num)
                    left_data = all_data
    
                    date_file.close()

        left_data = np.asarray(left_data)
        for i in range(len(left_data) // self.batch_size):
            train_data = left_data[i * self.batch_size: (i + 1) * self.batch_size]
            for j, compare_train in enumerate(compare_trains):
                compare_train.batch_train(train_data, train_data[:, self.fea_index['deal_type']].astype(np.int16))
        print("aucs = ", self.aucs)
        print("losses = ", self.losses)